from time import sleep
from mss.windows import MSS as mss
import multiprocessing
import pyautogui
import cv2
import numpy as np
from os import remove
from PIL import ImageGrab
import pygetwindow as pgw


class CaptureScreen(multiprocessing.Process):
    pause = True

    def __init__(self, video_path: str, fps: str):
        super().__init__()
        self.fps = fps
        self.video_path = video_path

    def capture_screen(self):

        SCREEN_SIZE = tuple(pyautogui.size())
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        print(self.video_path)
        out = cv2.VideoWriter(self.video_path, fourcc, self.fps, (SCREEN_SIZE))
        img = None
        print(f'Screen capture started')

        # MSS
        w, h = pyautogui.size()
        # (left, top, width, height)
        region = (0, 0, w, h)
        with mss() as sct:
            while self.pause:
                try:
                    try:
                        # = sct.grab(( 0,0,w,h) ) 
                        img = pyautogui.screenshot()
                        # img = ImageGrab.grab()
                        frame = np.array(img)
                        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
                        out.write(frame)
                    except IOError:
                        pass
                except KeyboardInterrupt:
                    break
            cv2.destroyAllWindows()
            out.release()

    def run(self):
        self.capture_screen()

    def terminate(self):
        if not self.process.is_alive():
            print('The process is not executing')
        else:
            super().terminate()
            sleep(3)
            print('Screen capture finished')


if __name__ == '__main__':
    c = CaptureScreen(video_path=r'E:\\game\\python-learn\\temp\\output\\output.avi', fps=20.0)
    c.start()
    sleep(15)
    c.terminate()
